# 声明：本代码仅供学习和研究目的使用。使用者应遵守以下原则：  
# 1. 不得用于任何商业用途。  
# 2. 使用时应遵守目标平台的使用条款和robots.txt规则。  
# 3. 不得进行大规模爬取或对平台造成运营干扰。  
# 4. 应合理控制请求频率，避免给目标平台带来不必要的负担。   
# 5. 不得用于任何非法或不当的用途。
#   
# 详细许可条款请参阅项目根目录下的LICENSE文件。  
# 使用本代码即表示您同意遵守上述原则和LICENSE中的所有条款。  


import asyncio
import copy
import json
import urllib.parse
from distutils.command.config import config
from typing import Any, Callable, Dict, Optional, List

import requests
from playwright.async_api import BrowserContext

from base.base_crawler import AbstractApiClient
from tools import utils
from var import request_keyword_var
from store import douyin as douyin_store
from .exception import *
from .field import *
from .help import *


class DOUYINClient(AbstractApiClient):
    def __init__(
            self,
            timeout=30,
            proxies=None,
            *,
            headers: Dict,
            playwright_page: Optional[Page],
            cookie_dict: Dict,
            max_notes_count: int,
            max_comments_count: int,
            max_fans_count: int,
            max_follower_count: int
    ):
        self.proxies = proxies
        self.timeout = timeout
        self.headers = headers
        self._host = "https://www.douyin.com"
        self.playwright_page = playwright_page
        self.cookie_dict = cookie_dict
        self.max_notes_count = max_notes_count
        self.max_comments_count = max_comments_count
        self.max_fans_count = max_fans_count
        self.max_follower_count = max_follower_count
        
        utils.logger.info(f"[DOUYINClient.__init__] Initialized with limits:")
        utils.logger.info(f"[DOUYINClient.__init__] - max_notes_count: {self.max_notes_count}")
        utils.logger.info(f"[DOUYINClient.__init__] - max_comments_count: {self.max_comments_count}")
        utils.logger.info(f"[DOUYINClient.__init__] - max_fans_count: {self.max_fans_count}")
        utils.logger.info(f"[DOUYINClient.__init__] - max_follower_count: {self.max_follower_count}")

    async def __process_req_params(
            self, uri: str, params: Optional[Dict] = None, headers: Optional[Dict] = None,
            request_method="GET"
    ):
        """处理请求参数"""
        if not params:
            return
            
        headers = headers or self.headers
        
        try:
            # 获取本地存储，最多重试3次
            max_retries = 3
            local_storage = None
            
            for retry in range(max_retries):
                try:
                    # 等待页面加载到可交互状态
                    await self.playwright_page.wait_for_load_state("domcontentloaded", timeout=10000)
                    local_storage = await self.playwright_page.evaluate("() => window.localStorage")
                    if local_storage:
                        break
                except Exception as e:
                    if retry == max_retries - 1:
                        raise Exception(f"获取本地存储失败: {str(e)}")
                    await asyncio.sleep(1)
            
            # 获取 msToken
            msToken = local_storage.get("xmst") if local_storage else None
            if not msToken:
                utils.logger.warning("[DOUYINClient] msToken不存在，尝试重新获取...")
                await asyncio.sleep(1)
                local_storage = await self.playwright_page.evaluate("() => window.localStorage")
                msToken = local_storage.get("xmst") if local_storage else None
            
            if not msToken:
                utils.logger.error("[DOUYINClient] 无法获取msToken")
            
            # 构建通用参数
            common_params = {
                "device_platform": "webapp",
                "aid": "6383",
                "channel": "channel_pc_web",
                "version_code": "190600",
                "version_name": "19.6.0",
                "cookie_enabled": "true",
                "screen_width": "1920",
                "screen_height": "1080",
                "browser_language": "zh-CN",
                "browser_platform": "MacIntel",
                "browser_name": "Chrome",
                "browser_version": "120.0.0.0",
                "browser_online": "true",
                "engine_name": "Blink",
                "engine_version": "120.0.0.0",
                "os_name": "Mac OS",
                "os_version": "10.15.7",
                "cpu_core_num": "8",
                "device_memory": "8",
                "platform": "PC",
                "downlink": "10",
                "effective_type": "4g",
                "round_trip_time": "50",
                "webid": get_web_id(),
                "msToken": msToken,
            }
            
            params.update(common_params)
            query_string = urllib.parse.urlencode(params)
            
            # 获取 a_bogus
            for i in range(max_retries):
                try:
                    a_bogus = await get_a_bogus(
                        uri, 
                        query_string,
                        {} if request_method == "GET" else params,
                        headers["User-Agent"],
                        self.playwright_page
                    )
                    if a_bogus:
                        params["a_bogus"] = a_bogus
                        break
                except Exception as e:
                    if i == max_retries - 1:
                        raise Exception(f"获取a_bogus失败: {str(e)}")
                    await asyncio.sleep(1)
                    
        except Exception as e:
            utils.logger.error(f"[DOUYINClient] 处理请求参数失败: {str(e)}")
            raise

    async def request(self, method, url, **kwargs):
        max_retries = 3
        retry_delay = 2
        
        for i in range(max_retries):
            try:
                # 添加代理和超时设置
                if self.proxies:
                    kwargs['proxies'] = self.proxies
                kwargs['timeout'] = self.timeout
                
                response = requests.request(method, url, **kwargs)
                
                # 检查响应状态
                if response.status_code == 429:  # Too Many Requests
                    utils.logger.warning(f"[DOUYINClient.request] Rate limited, waiting {retry_delay} seconds")
                    await asyncio.sleep(retry_delay * (i + 1))
                    continue
                    
                if response.status_code != 200:
                    utils.logger.error(f"[DOUYINClient.request] HTTP {response.status_code}: {response.text}")
                    raise Exception(f"HTTP {response.status_code}")
                    
                if response.text == "" or response.text == "blocked":
                    utils.logger.error(f"[DOUYINClient.request] Request blocked: {response.text}")
                    raise Exception("Account blocked")
                    
                return response.json()
                
            except Exception as e:
                if i == max_retries - 1:
                    utils.logger.error(f"[DOUYINClient.request] Request failed after {max_retries} retries: {str(e)}")
                    raise DataFetchError(f"Request failed: {str(e)}, Response: {response.text if 'response' in locals() else 'No response'}")
                utils.logger.warning(f"[DOUYINClient.request] Request failed (attempt {i+1}/{max_retries}): {str(e)}")
                await asyncio.sleep(retry_delay * (i + 1))

    async def get(self, uri: str, params: Optional[Dict] = None, headers: Optional[Dict] = None):
        """
        GET请求
        """
        await self.__process_req_params(uri, params, headers)
        headers = headers or self.headers
        return await self.request(method="GET", url=f"{self._host}{uri}", params=params, headers=headers)

    async def post(self, uri: str, data: dict, headers: Optional[Dict] = None):
        await self.__process_req_params(uri, data, headers)
        headers = headers or self.headers
        return await self.request(method="POST", url=f"{self._host}{uri}", data=data, headers=headers)

    async def pong(self, page: Page) -> bool:
        """检查登录状态
        
        Args:
            page: playwright page对象
            
        Returns:
            bool: 是否已登录
        """
        try:
            # 从页面获取登录状态
            local_storage = await page.evaluate("() => window.localStorage")
            if local_storage.get("HasUserLogin", "") == "1":
                return True

            # 从cookie检查登录状态
            cookies = await page.context.cookies()
            _, cookie_dict = utils.convert_cookies(cookies)
            return cookie_dict.get("LOGIN_STATUS") == "1"
            
        except Exception as e:
            utils.logger.error(f"[DOUYINClient.pong] 检查登录状态失败: {str(e)}")
            return False

    async def update_cookies(self, browser_context: BrowserContext):
        cookie_str, cookie_dict = utils.convert_cookies(await browser_context.cookies())
        self.headers["Cookie"] = cookie_str
        self.cookie_dict = cookie_dict

    async def search_info_by_keyword(
            self,
            keyword: str,
            offset: int = 0,
            search_channel: SearchChannelType = SearchChannelType.GENERAL,
            sort_type: SearchSortType = SearchSortType.GENERAL,
            publish_time: PublishTimeType = PublishTimeType.UNLIMITED,
            search_id: str = ""
    ):
        """抖音网页搜索 API"""
        max_retries = 3
        retry_delay = 2
        
        for retry in range(max_retries):
            try:
                # 确保已登录
                if not await self.pong(self.playwright_page):
                    utils.logger.warning("[DOUYINClient] 检测到未登录状态，尝试重新获取登录状态...")
                    await asyncio.sleep(2)
                    if not await self.pong(self.playwright_page):
                        raise Exception("登录状态校验失败")
                
                query_params = {
                    'search_channel': search_channel.value,
                    'enable_history': '1',
                    'keyword': keyword,
                    'search_source': 'tab_search',
                    'query_correct_type': '1',
                    'is_filter_search': '0',
                    'from_group_id': '7378810571505847586',
                    'offset': offset,
                    'count': '15',
                    'need_filter_settings': '1',
                    'list_type': 'multi',
                    'search_id': search_id,
                }
                
                if sort_type.value != SearchSortType.GENERAL.value or publish_time.value != PublishTimeType.UNLIMITED.value:
                    query_params["filter_selected"] = json.dumps({
                        "sort_type": str(sort_type.value),
                        "publish_time": str(publish_time.value)
                    })
                    query_params["is_filter_search"] = 1
                    query_params["search_source"] = "tab_search"
                
                # 构建 referer URL
                referer_url = f"https://www.douyin.com/search/{keyword}?aid=f594bbd9-a0e2-4651-9319-ebe3cb6298c1&type=general"
                headers = copy.copy(self.headers)
                headers["Referer"] = urllib.parse.quote(referer_url, safe=':/')
                
                response = await self.get("/aweme/v1/web/general/search/single/", query_params, headers=headers)
                
                if not response:
                    raise Exception("搜索响应为空")
                    
                if "status_code" in response and response["status_code"] != 0:
                    raise Exception(f"搜索请求失败: {response.get('status_msg', '未知错误')}")
                    
                return response
                
            except Exception as e:
                utils.logger.error(f"[DOUYINClient] 搜索失败 (尝试 {retry + 1}/{max_retries}): {str(e)}")
                if retry < max_retries - 1:
                    await asyncio.sleep(retry_delay * (retry + 1))
                    continue
                raise DataFetchError(f"搜索失败: {str(e)}")

    async def get_video_by_id(self, aweme_id: str) -> Any:
        """
        DouYin Video Detail API
        :param aweme_id:
        :return:
        """
        params = {
            "aweme_id": aweme_id
        }
        headers = copy.copy(self.headers)
        del headers["Origin"]
        res = await self.get("/aweme/v1/web/aweme/detail/", params, headers)
        return res.get("aweme_detail", {})

    async def get_aweme_comments(self, aweme_id: str, cursor: int = 0):
        """get note comments

        """
        uri = "/aweme/v1/web/comment/list/"
        params = {
            "aweme_id": aweme_id,
            "cursor": cursor,
            "count": 20,
            "item_type": 0
        }
        keywords = request_keyword_var.get()
        referer_url = "https://www.douyin.com/search/" + keywords + '?aid=3a3cec5a-9e27-4040-b6aa-ef548c2c1138&publish_time=0&sort_type=0&source=search_history&type=general'
        headers = copy.copy(self.headers)
        headers["Referer"] = urllib.parse.quote(referer_url, safe=':/')
        return await self.get(uri, params)

    async def get_sub_comments(self, comment_id: str, cursor: int = 0):
        """
            获取子评论
        """
        uri = "/aweme/v1/web/comment/list/reply/"
        params = {
            'comment_id': comment_id,
            "cursor": cursor,
            "count": 20,
            "item_type": 0,
        }
        keywords = request_keyword_var.get()
        referer_url = "https://www.douyin.com/search/" + keywords + '?aid=3a3cec5a-9e27-4040-b6aa-ef548c2c1138&publish_time=0&sort_type=0&source=search_history&type=general'
        headers = copy.copy(self.headers)
        headers["Referer"] = urllib.parse.quote(referer_url, safe=':/')
        return await self.get(uri, params)

    async def get_aweme_all_comments(
            self,
            aweme_id: str,
            crawl_interval: float = 1.0,
            is_fetch_sub_comments=False,
            callback: Optional[Callable] = None,
            max_count: int = None,
    ):
        """
        获取帖子的所有评论，包括子评论
        :param aweme_id: 帖子ID
        :param crawl_interval: 抓取间隔
        :param is_fetch_sub_comments: 是否抓取子评论
        :param callback: 回调函数，用于处理抓取到的评论
        :param max_count: 一次帖子爬取的最大评论数量，由前端传入
        :return: 评论列表
        """
        result = []
        comments_cursor = 0
        
        try:
            # 如果没有指定max_count，使用实例的max_comments_count
            if max_count is None:
                max_count = self.max_comments_count
                
            utils.logger.info(f"[DOUYINClient.get_aweme_all_comments] Getting comments for note {aweme_id}, max_count: {max_count}")
            
            while len(result) < max_count:
                comments_res = await self.get_aweme_comments(aweme_id, comments_cursor)
                if not comments_res or "comments" not in comments_res:
                    break
                    
                comments = comments_res.get("comments", [])
                if not comments:
                    break
                    
                # 确保不超过最大评论数
                remaining_count = max_count - len(result)
                if len(comments) > remaining_count:
                    comments = comments[:remaining_count]
                    
                result.extend(comments)
                utils.logger.info(f"[DOUYINClient.get_aweme_all_comments] Got {len(result)}/{max_count} comments for note {aweme_id}")
                
                if len(result) >= max_count:
                    break
                    
                # 检查是否有更多评论
                has_more = comments_res.get("has_more", 0)
                if not has_more:
                    break
                    
                comments_cursor = comments_res.get("cursor", 0)
                await asyncio.sleep(crawl_interval)
                
            return result
            
        except Exception as e:
            utils.logger.error(f"[DOUYINClient.get_aweme_all_comments] Error getting comments for note {aweme_id}: {str(e)}")
            return result

    async def get_user_info(self, sec_user_id: str):
        uri = "/aweme/v1/web/user/profile/other/"
        params = {
            "sec_user_id": sec_user_id,
            "publish_video_strategy_type": 2,
            "personal_center_strategy": 1,
        }
        return await self.get(uri, params)

    async def get_user_aweme_posts(self, sec_user_id: str, max_cursor: str = "") -> Dict:
        uri = "/aweme/v1/web/aweme/post/"
        params = {
            "sec_user_id": sec_user_id,
            "count": 18,
            "max_cursor": max_cursor,
            "locate_query": "false",
            "publish_video_strategy_type": 2,
            'verifyFp': 'verify_lx901cuk_K7kaK4dK_bn2E_4dgk_BxAA_E0XS1VtUi130',
            'fp': 'verify_lx901cuk_K7kaK4dK_bn2E_4dgk_BxAA_E0XS1VtUi130'
        }
        return await self.get(uri, params)

    async def get_all_user_aweme_posts(self, sec_user_id: str, callback: Optional[Callable] = None):
        posts_has_more = 1
        max_cursor = ""
        result = []
        while posts_has_more == 1:
            aweme_post_res = await self.get_user_aweme_posts(sec_user_id, max_cursor)
            posts_has_more = aweme_post_res.get("has_more", 0)
            max_cursor = aweme_post_res.get("max_cursor")
            aweme_list = aweme_post_res.get("aweme_list") if aweme_post_res.get("aweme_list") else []
            utils.logger.info(
                f"[DOUYINClient.get_all_user_aweme_posts] got sec_user_id:{sec_user_id} video len : {len(aweme_list)}")
            if callback:
                await callback(aweme_list)
            result.extend(aweme_list)
        return result


    async def get_creator_fans_get(self, sec_user_id: str, user_id: str, minTime: int) -> Dict:

        source_type = 2 if minTime == 0 else 1
        uri = "/aweme/v1/web/user/follower/list/"
        params = {
            "sec_user_id": sec_user_id,
            "offset": 0,
            "count": 20,
            "source_type": source_type,
            "pc_client_type": 1,
            "webid": "7447458789911135759",
            "max_time": minTime,
            "min_time": 0,
            'verifyFp': 'verify_m730p6yt_hb5alTSo_BEGr_40iM_9rPA_Vjnv4zg1dMRG',
            'fp': 'verify_m730p6yt_hb5alTSo_BEGr_40iM_9rPA_Vjnv4zg1dMRG',
            'gps_access': 0,
            'address_book_access': 0,
            "device_platform": "webapp",
            "aid": "6383",
            "user_id": user_id,
            "channel": "channel_pc_web",

        }
        return await self.get(uri, params)



    async def get_creator_fans(self,
                             sec_user_id: str,
                             user_id: str,
                             callback: Optional[Callable] = None,
                             max_count: int = 100) -> List[Dict]:
        """获取用户粉丝列表
        
        Args:
            sec_user_id: 用户的sec_user_id
            user_id: 用户ID
            callback: 回调函数
            max_count: 最大获取数量
            
        Returns:
            List[Dict]: 粉丝列表
        """
        fans_list = []
        min_time = 0
        
        try:
            # 获取初始 min_time
            initial_res = await self.get_creator_fans_get(sec_user_id, user_id, min_time)
            min_time = initial_res.get("min_time", 0)
            
            while len(fans_list) < max_count:
                res = await self.get_creator_fans_get(sec_user_id, user_id, min_time)
                if not res or "followers" not in res:
                    break
                    
                fans = res.get("followers", [])
                if not fans:
                    break
                    
                fans_list.extend(fans)
                
                if len(fans_list) >= max_count:
                    fans_list = fans_list[:max_count]
                    break
                    
                has_more = res.get("has_more", 0)
                if not has_more:
                    break
                    
                min_time = res.get("min_time", 0)
                await asyncio.sleep(1.0)  # 避免请求过快
                
            return fans_list
            
        except Exception as e:
            utils.logger.error(f"[DOUYINClient] 获取粉丝列表失败: {str(e)}")
            return []



    async def get_creator_followers_get(self, sec_user_id: str, user_id: str, offset: int) -> Dict:

        uri = "/aweme/v1/web/user/following/list"
        params = {
            "sec_user_id": sec_user_id,
            "offset": offset,
            "count": 20,
            "is_top": 1,
            "source_type": 4,
            "pc_client_type": 1,
            "webid": "7447458789911135759",
            "max_time": 0,
            "min_time": 0,
            'verifyFp': 'verify_m730p6yt_hb5alTSo_BEGr_40iM_9rPA_Vjnv4zg1dMRG',
            'fp': 'verify_m730p6yt_hb5alTSo_BEGr_40iM_9rPA_Vjnv4zg1dMRG',
            'gps_access': 0,
            'address_book_access': 0,
            "device_platform": "webapp",
            "aid": "6383",
            "user_id": user_id,
            "channel": "channel_pc_web",

        }
        return await self.get(uri, params)


    async def get_creator_followers(self,
                                   sec_user_id: str,
                                   user_id: str,
                                   callback: Optional[Callable] = None,
                                   max_count: int = 100) -> List[Dict]:
        """获取用户关注列表
        
        Args:
            sec_user_id: 用户的sec_user_id
            user_id: 用户ID
            callback: 回调函数
            max_count: 最大获取数量
            
        Returns:
            List[Dict]: 关注列表
        """
        followers_list = []
        offset = 0
        
        try:
            while len(followers_list) < max_count:
                res = await self.get_creator_followers_get(sec_user_id, user_id, offset)
                if not res or "followings" not in res:
                    break
                    
                followers = res.get("followings", [])
                if not followers:
                    break
                    
                followers_list.extend(followers)
                
                if len(followers_list) >= max_count:
                    followers_list = followers_list[:max_count]
                    break
                    
                has_more = res.get("has_more", 0)
                if not has_more:
                    break
                    
                offset = res.get("offset", 0)
                await asyncio.sleep(1.0)  # 避免请求过快
                
            return followers_list
            
        except Exception as e:
            utils.logger.error(f"[DOUYINClient] 获取关注列表失败: {str(e)}")
            return []

    async def get_creator_notes(self, creator_id: str, need_comments: bool = False) -> List[Dict]:
        """Get notes from a creator
        Args:
            creator_id: Creator's ID
            need_comments: 是否需要获取评论，默认为False
        Returns:
            List of notes
        """
        try:
            utils.logger.info(f"[DOUYINClient.get_creator_notes] Getting notes for creator: {creator_id}")
            
            # 获取创作者信息并存储
            creator_info = await self.get_user_info(creator_id)
            if creator_info and "user" in creator_info:
                await douyin_store.save_creator(creator_id, creator_info)
                utils.logger.info(f"[DOUYINClient.get_creator_notes] Saved creator info for {creator_id}")
            
            # 获取所有帖子
            creator_notes = await self.get_all_user_aweme_posts(
                sec_user_id=creator_id,
                callback=None
            )
            
            # 随机选择指定数量的帖子
            if creator_notes:
                import random
                if len(creator_notes) > self.max_notes_count:
                    notes = random.sample(creator_notes, self.max_notes_count)
                else:
                    notes = creator_notes
                    
                utils.logger.info(f"[DOUYINClient.get_creator_notes] Got {len(notes)} notes from creator {creator_id}")
                
                # 存储帖子信息，根据need_comments决定是否获取评论
                for note in notes:
                    try:
                        await douyin_store.update_douyin_aweme(note)
                        utils.logger.info(f"[DOUYINClient.get_creator_notes] Saved note {note.get('aweme_id', 'unknown')}")
                        
                        if need_comments and "aweme_id" in note:
                            utils.logger.info(f"[DOUYINClient.get_creator_notes] Getting comments for note {note['aweme_id']}")
                            comments = await self.get_aweme_all_comments(
                                aweme_id=note["aweme_id"],
                                max_count=self.max_comments_count,
                                is_fetch_sub_comments=False
                            )
                            if comments:
                                for comment in comments:
                                    await douyin_store.update_dy_aweme_comment(note["aweme_id"], comment)
                                utils.logger.info(f"[DOUYINClient.get_creator_notes] Saved {len(comments)} comments for note {note['aweme_id']}")
                    except Exception as e:
                        utils.logger.error(f"[DOUYINClient.get_creator_notes] Error processing note {note.get('aweme_id', 'unknown')}: {str(e)}")
                        continue
                
                # 获取粉丝和关注者数据
                if creator_info and "user" in creator_info:
                    user_id = creator_info["user"].get("uid")
                    if user_id:
                        try:
                            # 获取并存储粉丝关系
                            utils.logger.info(f"[DOUYINClient.get_creator_notes] Getting fans for creator {creator_id}")
                            fans = await self.get_creator_fans(creator_id, user_id, max_count=self.max_fans_count)
                            for fan in fans:
                                relation_data = {
                                    "user_id": creator_id,
                                    "r_user_id": fan.get("sec_uid", ""),
                                    "r_type": "0"  # 0表示粉丝关系
                                }
                                await douyin_store.save_relation_user(relation_data)
                                
                            # 获取并存储关注关系
                            utils.logger.info(f"[DOUYINClient.get_creator_notes] Getting followers for creator {creator_id}")
                            followers = await self.get_creator_followers(creator_id, user_id, max_count=self.max_follower_count)
                            for follower in followers:
                                relation_data = {
                                    "user_id": creator_id,
                                    "r_user_id": follower.get("sec_uid", ""),
                                    "r_type": "1"  # 1表示关注关系
                                }
                                await douyin_store.save_relation_user(relation_data)
                        except Exception as e:
                            utils.logger.error(f"[DOUYINClient.get_creator_notes] Error getting relations: {str(e)}")
                
                return notes
            else:
                utils.logger.warning(f"[DOUYINClient.get_creator_notes] No notes found for creator {creator_id}")
                return []
                
        except Exception as e:
            utils.logger.error(f"[DOUYINClient.get_creator_notes] Error getting notes for creator {creator_id}: {str(e)}")
            import traceback
            utils.logger.error(f"[DOUYINClient.get_creator_notes] Error traceback: {traceback.format_exc()}")
            return []
